# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
from typing import List, Union

from coin.base.param_util import to_list
from coin.exchange.base.kr_rest.product import NormProduct
from coin.feed.fastfeed.proto_feed_handler import create_proto_feed_handler
from coin.strategy.mm.feed_subsystems import get_feed_subsystem_cls
from coin.strategy.mm.subscription import FeedSubscriptionRequest, get_feed_sub_request


class _Feeder(object):
  def __init__(self, subscription_key, feed_checker_config=None, fs_config=None, logger=None):
    self._logger = logger or logging.getLogger('_Feeder')
    self._subscription_key = subscription_key
    self._products = set()
    self._subscription = {}
    self._feed_subsystem = None
    self._feed_checker_config = feed_checker_config
    self._fs_config = fs_config

    self._fastfeed_book_reset_callbacks = []
    self._fastfeed_handler = None

  def subscribe_event(self, event_type, key, callback):
    assert not self.prepared
    event_key = (event_type, key)
    if event_key not in self._subscription:
      self._subscription[event_key] = []

    if callback is not None:
      self._subscription[event_key].append(callback)

    if event_type in ('book', 'trade'):
      self._products.add(key)

  def prepare(self):
    assert not self.prepared
    feed_subsystem_cls = get_feed_subsystem_cls(self._subscription_key, fs_config=self._fs_config)

    self._feed_subsystem = feed_subsystem_cls(list(self._products),
                                              on_book_init_callback=self.on_book_reset,
                                              feed_checker_config=self._feed_checker_config,
                                              fs_config=self._fs_config)

  @property
  def prepared(self):
    return self._feed_subsystem is not None

  @property
  def available(self):
    if self._feed_subsystem is None:
      return False
    return self._feed_subsystem.available

  def start(self, ioloop, queue, flow_sub):
    self._feed_subsystem.start(ioloop, queue, flow_sub)

  def restart(self):
    self._feed_subsystem.restart()

  def on_book_reset(self, sub_req, book_builder):
    self._logger.info('Feed initialized: %s', sub_req)

    for event_key, callbacks in self._subscription.items():
      for callback in callbacks:
        book_builder.subscribe_event(event_key[0], event_key[1], callback)

    self.invoke_fastfeed_book_reset_callbacks(book_builder)

  def add_fastfeed_book_reset_callback(self, callback):
    self._fastfeed_book_reset_callbacks += to_list(callback)

  def invoke_fastfeed_book_reset_callbacks(self, book_builder):
    if not self._fastfeed_book_reset_callbacks:
      return

    # TODO(jaewon): Support other types such as instrument info.
    self._fastfeed_handler = create_proto_feed_handler(self._subscription_key, book_builder)
    products = sorted(self._products)
    self._fastfeed_handler.subscribe_book(products, None)
    self._fastfeed_handler.subscribe_trade(products, None)

    from cc.coin1.fastfeed.pyfastfeed import _FeedCacheKey
    feed_cache_key = _FeedCacheKey()
    feed_cache_key.sub_req = self._subscription_key.to_str()

    for callback in self._fastfeed_book_reset_callbacks:
      callback(self._subscription_key, self._fastfeed_handler, feed_cache_key=feed_cache_key)


class MultiFeedSubsystem(object):
  def __init__(self,
               *,
               feed_checker_config=None,
               feed_checker_config_map=None,
               fs_config=None,
               logger=None):
    self._logger = logger or logging.getLogger('MultiFeedSubsystem')
    self._feeders = {}
    self._feed_checker_config = feed_checker_config
    self._feed_checker_config_map = feed_checker_config_map
    self._fs_config = fs_config

  def set_price_limit_error_callback(self, subscription_key, callback):
    assert subscription_key in self._feeders, (subscription_key, self._feeders)
    feeder = self._feeders.get(subscription_key)
    feeder._feed_subsystem.set_price_limit_error_callback(callback)

  def _get_feed_checker_config(self, subscription_key):
    if (self._feed_checker_config_map is not None
        and subscription_key in self._feed_checker_config_map):
      return self._feed_checker_config_map[subscription_key]
    else:
      return self._feed_checker_config

  def subscribe_event(self, subscription_key, event_type, keys, callback):
    if subscription_key not in self._feeders:
      self._feeders[subscription_key] = _Feeder(
          subscription_key,
          feed_checker_config=self._get_feed_checker_config(subscription_key),
          fs_config=self._fs_config,
          logger=self._logger)
    for key in to_list(keys):
      self._feeders[subscription_key].subscribe_event(event_type, key, callback)

  def subscribe_book(self,
                     feed_sub_request_or_products: Union[FeedSubscriptionRequest,
                                                         List[NormProduct]],
                     callback):
    feed_sub_request = get_feed_sub_request(feed_sub_request_or_products)
    for sub_key, product_list in feed_sub_request.gen_iter():
      for product in product_list:
        self.subscribe_event(sub_key, 'book', product, callback)

  def subscribe_trade(self,
                      feed_sub_request_or_products: Union[FeedSubscriptionRequest,
                                                          List[NormProduct]],
                      callback):
    feed_sub_request = get_feed_sub_request(feed_sub_request_or_products)
    for sub_key, product_list in feed_sub_request.gen_iter():
      for product in product_list:
        self.subscribe_event(sub_key, 'trade', product, callback)

  def subscribe_feed_status(self,
                            feed_sub_request_or_products: Union[FeedSubscriptionRequest,
                                                                List[NormProduct]],
                            callback):
    feed_sub_request = get_feed_sub_request(feed_sub_request_or_products)
    for sub_key, _ in feed_sub_request.gen_iter():
      self.subscribe_event(sub_key, 'feed_status', 'feed_status', callback)

  def subscribe_fastfeed(self,
                         feed_sub_request_or_products: Union[FeedSubscriptionRequest,
                                                             List[NormProduct]],
                         fastfeed_book_reset_callback):
    feed_sub_request = get_feed_sub_request(feed_sub_request_or_products)
    self.subscribe_book(feed_sub_request, None)
    self.subscribe_trade(feed_sub_request, None)
    for sub_key, _ in feed_sub_request.gen_iter():
      self._feeders[sub_key].add_fastfeed_book_reset_callback(fastfeed_book_reset_callback)

  def start(self, ioloop, queue, flow_sub, check_interval=30):
    for feeder in self._feeders.values():
      feeder.prepare()
      feeder.start(ioloop, queue, flow_sub)
      feeder._feed_subsystem._book_builder.start_periodic_check(ioloop, check_interval)

  def restart(self):
    for feeder in self._feeders:
      feeder.restart()

  @property
  def available(self):
    for feeder in self._feeders.values():
      if not feeder.available:
        return False
    return True
